import itertools
import pandas as pd

# 步骤 1：读取数据
data = [
    ['Milk', 'Bread', 'Butter'],
    ['Milk', 'Bread'],
    ['Milk', 'Butter'],
    ['Bread', 'Butter'],
    ['Milk', 'Bread', 'Butter', 'Cheese'],
    ['Bread', 'Butter', 'Cheese']
]


# 步骤 2：生成候选项集函数
def create_candidates(frequent_itemsets, k):
    """
    根据频繁项集生成新的候选项集。

    参数:
        frequent_itemsets (list): 上一轮迭代中找到的频繁项集列表。
        k (int): 当前要生成的候选项集的大小（项数）。

    返回:
        list: 新的候选项集列表。
    """
    if k == 1:
        # 如果是第一次生成候选项集，则直接从所有事务中提取唯一项
        items = set(itertools.chain.from_iterable(data))
        return [(item,) for item in items]
    else:
        # 合并上一轮频繁项集中的项，生成新的候选项集
        candidates = set()
        for i in range(len(frequent_itemsets)):
            for j in range(i + 1, len(frequent_itemsets)):
                itemset_i = set(frequent_itemsets[i])
                itemset_j = set(frequent_itemsets[j])
                if len(itemset_i.symmetric_difference(itemset_j)) == 2:
                    candidate = tuple(sorted(itemset_i.union(itemset_j)))
                    candidates.add(candidate)
        return list(candidates)


# 步骤 3：计算支持度函数
def compute_support(data, candidates):
    """
    计算每个候选项集的支持度。

    参数:
        data (list): 所有交易记录。
        candidates (list): 候选项集列表。

    返回:
        dict: 每个候选项集及其对应的支持度。
    """
    support = {}
    total_transactions = len(data)
    for candidate in candidates:
        count = sum(1 for transaction in data if set(candidate).issubset(set(transaction)))
        support[candidate] = count / total_transactions
    return support


# 步骤 4：Apriori算法主函数
def apriori(data, min_support):
    """
    实现Apriori算法，找出满足最小支持度的所有频繁项集。

    参数:
        data (list): 所有交易记录。
        min_support (float): 最小支持度阈值。

    返回:
        list: 每轮迭代中找到的频繁项集及其支持度。
    """
    frequent_itemsets = []  # 存储所有频繁项集
    k = 1  # 项集大小起始值

    # 初始候选项集为单个商品的集合
    candidates = create_candidates([], k)  # 注意这里传入空列表，因为首次调用时没有之前的频繁项集
    while candidates:
        support = compute_support(data, candidates)

        # 过滤出满足最小支持度的频繁项集
        frequent_itemsets_k = {item: sup for item, sup in support.items() if sup >= min_support}
        if not frequent_itemsets_k:
            break  # 如果没有频繁项集，停止迭代

        frequent_itemsets.append(frequent_itemsets_k)
        candidates = create_candidates(list(frequent_itemsets_k.keys()), k + 1)
        k += 1

    return frequent_itemsets


# 步骤 5：调用Apriori算法并输出结果
min_support = 0.5
frequent_itemsets = apriori(data, min_support)

# 输出结果
for k, itemsets in enumerate(frequent_itemsets, start=1):
    print(f"频繁{k}-项集:")
    for itemset, support in itemsets.items():
        print(f"  {itemset}: 支持度 = {support:.2f}")